package com.example.spring.ai.config;

import com.example.spring.ai.tools.DateTimeTools;
import com.knuddels.jtokkit.api.EncodingType;
import org.springframework.ai.autoconfigure.vectorstore.redis.RedisVectorStoreProperties;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.memory.InMemoryChatMemory;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.embedding.TokenCountBatchingStrategy;
import org.springframework.ai.openai.OpenAiChatModel;
import org.springframework.ai.tool.ToolCallbacks;
import org.springframework.ai.vectorstore.SimpleVectorStore;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.vectorstore.redis.RedisVectorStore;
import org.springframework.boot.autoconfigure.data.redis.RedisProperties;
import org.springframework.boot.context.properties.EnableConfigurationProperties;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import redis.clients.jedis.JedisPooled;

@EnableConfigurationProperties({RedisProperties.class, RedisVectorStoreProperties.class})
@Configuration
public class MyAiConfig {
    @Bean
    public ChatClient openAiClient(OpenAiChatModel model, ChatMemory chatMemory) {
        return ChatClient.builder(model)
                .defaultSystem("""
                        你是一个严谨的聊天机器人,
                        要根据已有数据进行回答,不能随便捏造数据,
                        并且用{voice}的语气回复用户
                        """
                )
                .defaultAdvisors(
                        new MessageChatMemoryAdvisor(chatMemory)
                )
                .build();
    }

    @Bean
    public ChatClient chatClientWithTools(OpenAiChatModel model) {
        return ChatClient.builder(model)
//                .defaultTools(new DateTimeTools())
                .defaultTools(ToolCallbacks.from(new DateTimeTools()))
                .build();
    }

    @Bean
    public ChatMemory chatMemory() {
        return new InMemoryChatMemory();
    }

    @Bean
    public VectorStore simpleVectorStore(EmbeddingModel embeddingModel) {
        return SimpleVectorStore.builder(embeddingModel)
                .batchingStrategy(new TokenCountBatchingStrategy(
                        EncodingType.CL100K_BASE,
                        1000,
                        0.1
                ))
                .build();
    }

    @Bean
    public RedisVectorStore redisVectorStore(JedisPooled jedis,
                                             EmbeddingModel embeddingModel,
                                             RedisVectorStoreProperties redisVectorStoreProperties) {
        return RedisVectorStore.builder(jedis, embeddingModel)
                .prefix(redisVectorStoreProperties.getPrefix())
                .indexName(redisVectorStoreProperties.getIndex())
                .metadataFields(RedisVectorStore.MetadataField.numeric("sku_id"),
                        RedisVectorStore.MetadataField.numeric("price"))
                .initializeSchema(redisVectorStoreProperties.isInitializeSchema())
                .build();
    }

    @Bean
    public JedisPooled jedisPooled(RedisProperties redisProperties) {
        return new JedisPooled(redisProperties.getHost(), redisProperties.getPort());
    }
}
